import pandas as pd
import numpy as np
from scipy.optimize import minimize
try:
    from gurobi import *
except:
    from gurobipy import *
import time

pd.options.display.max_columns = 200
pd.options.display.max_rows = 1000
pd.set_option('max_info_columns', 200)
pd.set_option('expand_frame_repr', False)
pd.set_option('expand_frame_repr', True)
pd.set_option('max_colwidth',1000)
pd.set_option('display.width',None)
pd.set_option('display.float_format', lambda x: '%.3f' % x)

#SET TO DIRECTORY HOUSING output/acr_weight_input.dta
os.chdir("") 

P_list = {}

data = pd.read_stata("acr_weight_input.dta")
data = data.loc[data.time.notnull() & (data.time <= 20)]
P_list[0] = [1] + list(data.fitted25.values)
P_list[1] = [1] + list(data.fitted75.values)

nz=len(P_list)
ncells = len(P_list[0])

Pr = {}
for z in range(nz):
    for c in range(len(P_list[z])):
        Pr[(z,c)]=P_list[z][c]
        
# start gurobi model
m = Model("LP")

# define types that go from j to k to l to m (depending on how many z) as they move from judge 0 to 1 to 2 to N
if nz==2:
    mass_type = {(j,k)     : .5 for j in range(ncells) for k in range(j,ncells)}
elif nz==3:
    mass_type = {(j,k,l)   : .5 for j in range(ncells) for k in range(j,ncells) for l in range(k,ncells)}
elif nz==4:
    mass_type = {(j,k,l,m) : .5 for j in range(ncells) for k in range(j,ncells) for l in range(k,ncells) for m in range(l,ncells)}
else:
    BREAK

mass_type  = m.addVars(mass_type,lb=0,ub=1,name="mass_type")

# adding up 
m.addLConstr(sum([mass_type[key] for key in mass_type.keys()])==1)

# fit marginal distribution
for c in range(ncells):
    for z in range(nz):
        m.addLConstr( Pr[(z,c)] == quicksum([mass_type[key] for key in mass_type.keys() if key[z]>=c]))

# find the largest and smallest number of compliers
Obj = quicksum([mass_type[key] for key in mass_type.keys() if (key[0]!=key[nz-1])])

m.setObjective(Obj , GRB.MAXIMIZE) 
m.setParam("OutputFlag",0)    
m.optimize()
max_comp = m.ObjVal

m.setObjective(Obj , GRB.MINIMIZE) 
m.setParam("OutputFlag",0)    
m.optimize()
min_comp = m.ObjVal

# calculate upper and lower bound of extensive margin share
mass_types_cons = m.getAttr('x', mass_type) 
for thresh_bin in range(ncells):
    lb = sum([mass_types_cons[key] for key in mass_types_cons.keys() if key[0]<=thresh_bin and key[nz-1]>thresh_bin]) / max_comp
    ub = sum([mass_types_cons[key] for key in mass_types_cons.keys() if key[0]<=thresh_bin and key[nz-1]>thresh_bin]) / min_comp
    
    print('bounds')
    print(lb,ub)

